//
// ImageReader.cs
//
// Author:
//   Jb Evain (jbevain@gmail.com)
//
// Copyright (c) 2008 - 2011 Jb Evain
//
// Permission is hereby granted, free of charge, to any person obtaining
// a copy of this software and associated documentation files (the
// "Software"), to deal in the Software without restriction, including
// without limitation the rights to use, copy, modify, merge, publish,
// distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to
// the following conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
//

using System;
using System.IO;

using Mono.Cecil.Metadata;

using RVA = System.UInt32;

namespace Mono.Cecil.PE
{

	sealed class ImageReader : BinaryStreamReader
	{

		readonly Image image;

		DataDirectory cli;
		DataDirectory metadata;

		public ImageReader(Stream stream)
			: base(stream)
		{
			image = new Image();

			image.FileName = stream.GetFullyQualifiedName();
		}

		void MoveTo(DataDirectory directory)
		{
			BaseStream.Position = image.ResolveVirtualAddress(directory.VirtualAddress);
		}

		void MoveTo(uint position)
		{
			BaseStream.Position = position;
		}

		void ReadImage()
		{
			if(BaseStream.Length < 128)
				throw new BadImageFormatException();

			// - DOSHeader

			// PE					2
			// Start				58
			// Lfanew				4
			// End					64

			if(ReadUInt16() != 0x5a4d)
				throw new BadImageFormatException();

			Advance(58);

			MoveTo(ReadUInt32());

			if(ReadUInt32() != 0x00004550)
				throw new BadImageFormatException();

			// - PEFileHeader

			// Machine				2
			image.Architecture = ReadArchitecture();

			// NumberOfSections		2
			ushort sections = ReadUInt16();

			// TimeDateStamp		4
			// PointerToSymbolTable	4
			// NumberOfSymbols		4
			// OptionalHeaderSize	2
			Advance(14);

			// Characteristics		2
			ushort characteristics = ReadUInt16();

			ushort subsystem;
			ReadOptionalHeaders(out subsystem);
			ReadSections(sections);
			if(!cli.IsZero)
			{
				ReadCLIHeader();
				ReadMetadata();
			}
			image.Kind = GetModuleKind(characteristics, subsystem);

		}

		TargetArchitecture ReadArchitecture()
		{
			var machine = ReadUInt16();
			switch(machine)
			{
				case 0x014c:
					return TargetArchitecture.I386;
				case 0x8664:
					return TargetArchitecture.AMD64;
				case 0x0200:
					return TargetArchitecture.IA64;
				case 0x01c4:
					return TargetArchitecture.ARMv7;
			}

			throw new NotSupportedException(machine.ToString());
		}

		static ModuleKind GetModuleKind(ushort characteristics, ushort subsystem)
		{
			if((characteristics & 0x2000) != 0) // ImageCharacteristics.Dll
				return ModuleKind.Dll;

			if(subsystem == 0x2 || subsystem == 0x9) // SubSystem.WindowsGui || SubSystem.WindowsCeGui
				return ModuleKind.Windows;

			return ModuleKind.Console;
		}

		void ReadOptionalHeaders(out ushort subsystem)
		{
			// - PEOptionalHeader
			//   - StandardFieldsHeader

			// Magic				2
			bool pe64 = ReadUInt16() == 0x20b;
			if(pe64)
				throw new NotSupportedException();
			image.NtHeader = new NtHeader();

			//						pe32 || pe64

			// LMajor				1
			// LMinor				1
			// CodeSize				4
			// InitializedDataSize	4
			// UninitializedDataSize4
			Advance(14);
			image.NtHeader.EntryPoint = ReadUInt32();
			// EntryPointRVA		4
			// BaseOfCode			4
			// BaseOfData			4 || 0


			Advance(8);
			//   - NTSpecificFieldsHeader
			image.NtHeader.ImageBase = ReadUInt32(); // ImageBase			4 || 8



			// SectionAlignment		4
			// FileAlignement		4
			// OSMajor				2
			// OSMinor				2
			// UserMajor			2
			// UserMinor			2
			// SubSysMajor			2
			// SubSysMinor			2
			// Reserved				4
			Advance(24);
			image.NtHeader.ImageSize = ReadUInt32(); // ImageSize			4

			// HeaderSize			4
			// FileChecksum			4
			Advance(8);


			subsystem = ReadUInt16(); // SubSystem			2

			// DLLFlags				2
			// StackReserveSize		4 || 8
			// StackCommitSize		4 || 8
			// HeapReserveSize		4 || 8
			// HeapCommitSize		4 || 8
			// LoaderFlags			4
			// NumberOfDataDir		4

			//   - DataDirectoriesHeader
			Advance(pe64 ? 42 : 26);
			image.Export = ReadDataDirectory(); // ExportTable			8
			// ImportTable			8
			// ResourceTable		8
			// ExceptionTable		8
			// CertificateTable		8
			// BaseRelocationTable	8

			Advance(40);

			// Debug				8
			image.Debug = ReadDataDirectory();

			// Copyright			8
			// GlobalPtr			8
			// TLSTable				8
			// LoadConfigTable		8
			// BoundImport			8
			// IAT					8
			// DelayImportDescriptor8
			Advance(56);

			// CLIHeader			8
			cli = ReadDataDirectory();

			//if(cli.IsZero)
			//	throw new BadImageFormatException();

			// Reserved				8
			Advance(8);
		}

		string ReadAlignedString(int length)
		{
			int read = 0;
			var buffer = new char[length];
			while(read < length)
			{
				var current = ReadByte();
				if(current == 0)
					break;

				buffer[read++] = (char)current;
			}

			Advance(-1 + ((read + 4) & ~3) - read);

			return new string(buffer, 0, read);
		}

		string ReadZeroTerminatedString(int length)
		{
			int read = 0;
			var buffer = new char[length];
			var bytes = ReadBytes(length);
			while(read < length)
			{
				var current = bytes[read];
				if(current == 0)
					break;

				buffer[read++] = (char)current;
			}

			return new string(buffer, 0, read);
		}

		void ReadSections(ushort count)
		{
			var sections = new Section[count];

			for(int i = 0 ; i < count ; i++)
			{
				var section = new Section();

				// Name
				section.Name = ReadZeroTerminatedString(8);

				// VirtualSize		4
				section.VirtualSize = ReadUInt32();

				// VirtualAddress	4
				section.VirtualAddress = ReadUInt32();
				// SizeOfRawData	4
				section.SizeOfRawData = ReadUInt32();
				// PointerToRawData	4
				section.PointerToRawData = ReadUInt32();

				// PointerToRelocations		4
				// PointerToLineNumbers		4
				// NumberOfRelocations		2
				// NumberOfLineNumbers		2
				// Characteristics			4
				Advance(16);

				sections[i] = section;

				ReadSectionData(section);
			}

			image.Sections = sections;
		}

		void ReadSectionData(Section section)
		{
			var position = BaseStream.Position;

			MoveTo(section.PointerToRawData);

			var length = (int)section.SizeOfRawData;
			var data = new byte[length];
			int offset = 0, read;

			while((read = Read(data, offset, length - offset)) > 0)
				offset += read;

			section.Data = data;

			BaseStream.Position = position;
		}

		void ReadCLIHeader()
		{
			MoveTo(cli);

			// - CLIHeader

			// Cb						4
			// MajorRuntimeVersion		2
			// MinorRuntimeVersion		2
			Advance(8);

			// Metadata					8
			metadata = ReadDataDirectory();
			// Flags					4
			image.Attributes = (ModuleAttributes)ReadUInt32();
			// EntryPointToken			4
			image.EntryPointToken = ReadUInt32();
			// Resources				8
			image.Resources = ReadDataDirectory();
			// StrongNameSignature		8
			// CodeManagerTable			8
			// VTableFixups				8
			// ExportAddressTableJumps	8
			// ManagedNativeHeader		8
		}

		void ReadMetadata()
		{
			MoveTo(metadata);

			if(ReadUInt32() != 0x424a5342)
				throw new BadImageFormatException();

			// MajorVersion			2
			// MinorVersion			2
			// Reserved				4
			Advance(8);

			var version = ReadZeroTerminatedString(ReadInt32());
			image.Runtime = version.ParseRuntime();

			// Flags		2
			Advance(2);

			var streams = ReadUInt16();

			var section = image.GetSectionAt(metadata.VirtualAddress);
			if(section == null)
				throw new BadImageFormatException();

			image.MetadataSection = section;

			for(int i = 0 ; i < streams ; i++)
				ReadMetadataStream(section);

			if(image.TableHeap != null)
				ReadTableHeap();
		}

		void ReadMetadataStream(Section section)
		{
			// Offset		4
			uint start = metadata.VirtualAddress - section.VirtualAddress + ReadUInt32(); // relative to the section start

			// Size			4
			uint size = ReadUInt32();

			var name = ReadAlignedString(16);
			switch(name)
			{
				case "#~":
				case "#-":
					image.TableHeap = new TableHeap(section, start, size);
					break;
				case "#Strings":
					image.StringHeap = new StringHeap(section, start, size);
					break;
				case "#Blob":
					image.BlobHeap = new BlobHeap(section, start, size);
					break;
				case "#GUID":
					image.GuidHeap = new GuidHeap(section, start, size);
					break;
				case "#US":
					image.UserStringHeap = new UserStringHeap(section, start, size);
					break;
			}
		}

		void ReadTableHeap()
		{
			var heap = image.TableHeap;

			uint start = heap.Section.PointerToRawData;

			MoveTo(heap.Offset + start);

			// Reserved			4
			// MajorVersion		1
			// MinorVersion		1
			Advance(6);

			// HeapSizes		1
			var sizes = ReadByte();

			// Reserved2		1
			Advance(1);

			// Valid			8
			heap.Valid = ReadInt64();

			// Sorted			8
			heap.Sorted = ReadInt64();

			for(int i = 0 ; i < TableHeap.TableCount ; i++)
			{
				if(!heap.HasTable((Table)i))
					continue;

				heap.Tables[i].Length = ReadUInt32();
			}

			SetIndexSize(image.StringHeap, sizes, 0x1);
			SetIndexSize(image.GuidHeap, sizes, 0x2);
			SetIndexSize(image.BlobHeap, sizes, 0x4);

			ComputeTableInformations();
		}

		static void SetIndexSize(Heap heap, uint sizes, byte flag)
		{
			if(heap == null)
				return;

			heap.IndexSize = (sizes & flag) > 0 ? 4 : 2;
		}

		int GetTableIndexSize(Table table)
		{
			return image.GetTableIndexSize(table);
		}

		int GetCodedIndexSize(CodedIndex index)
		{
			return image.GetCodedIndexSize(index);
		}

		void ComputeTableInformations()
		{
			uint offset = (uint)BaseStream.Position - image.MetadataSection.PointerToRawData; // header

			int stridx_size = image.StringHeap.IndexSize;
			int blobidx_size = image.BlobHeap != null ? image.BlobHeap.IndexSize : 2;

			var heap = image.TableHeap;
			var tables = heap.Tables;

			for(int i = 0 ; i < TableHeap.TableCount ; i++)
			{
				var table = (Table)i;
				if(!heap.HasTable(table))
					continue;

				int size;
				switch(table)
				{
					case Table.Module:
						size = 2	// Generation
							+ stridx_size	// Name
							+ (image.GuidHeap.IndexSize * 3);	// Mvid, EncId, EncBaseId
						break;
					case Table.TypeRef:
						size = GetCodedIndexSize(CodedIndex.ResolutionScope)	// ResolutionScope
							+ (stridx_size * 2);	// Name, Namespace
						break;
					case Table.TypeDef:
						size = 4	// Flags
							+ (stridx_size * 2)	// Name, Namespace
							+ GetCodedIndexSize(CodedIndex.TypeDefOrRef)	// BaseType
							+ GetTableIndexSize(Table.Field)	// FieldList
							+ GetTableIndexSize(Table.Method);	// MethodList
						break;
					case Table.FieldPtr:
						size = GetTableIndexSize(Table.Field);	// Field
						break;
					case Table.Field:
						size = 2	// Flags
							+ stridx_size	// Name
							+ blobidx_size;	// Signature
						break;
					case Table.MethodPtr:
						size = GetTableIndexSize(Table.Method);	// Method
						break;
					case Table.Method:
						size = 8	// Rva 4, ImplFlags 2, Flags 2
							+ stridx_size	// Name
							+ blobidx_size	// Signature
							+ GetTableIndexSize(Table.Param); // ParamList
						break;
					case Table.ParamPtr:
						size = GetTableIndexSize(Table.Param); // Param
						break;
					case Table.Param:
						size = 4	// Flags 2, Sequence 2
							+ stridx_size;	// Name
						break;
					case Table.InterfaceImpl:
						size = GetTableIndexSize(Table.TypeDef)	// Class
							+ GetCodedIndexSize(CodedIndex.TypeDefOrRef);	// Interface
						break;
					case Table.MemberRef:
						size = GetCodedIndexSize(CodedIndex.MemberRefParent)	// Class
							+ stridx_size	// Name
							+ blobidx_size;	// Signature
						break;
					case Table.Constant:
						size = 2	// Type
							+ GetCodedIndexSize(CodedIndex.HasConstant)	// Parent
							+ blobidx_size;	// Value
						break;
					case Table.CustomAttribute:
						size = GetCodedIndexSize(CodedIndex.HasCustomAttribute)	// Parent
							+ GetCodedIndexSize(CodedIndex.CustomAttributeType)	// Type
							+ blobidx_size;	// Value
						break;
					case Table.FieldMarshal:
						size = GetCodedIndexSize(CodedIndex.HasFieldMarshal)	// Parent
							+ blobidx_size;	// NativeType
						break;
					case Table.DeclSecurity:
						size = 2	// Action
							+ GetCodedIndexSize(CodedIndex.HasDeclSecurity)	// Parent
							+ blobidx_size;	// PermissionSet
						break;
					case Table.ClassLayout:
						size = 6	// PackingSize 2, ClassSize 4
							+ GetTableIndexSize(Table.TypeDef);	// Parent
						break;
					case Table.FieldLayout:
						size = 4	// Offset
							+ GetTableIndexSize(Table.Field);	// Field
						break;
					case Table.StandAloneSig:
						size = blobidx_size;	// Signature
						break;
					case Table.EventMap:
						size = GetTableIndexSize(Table.TypeDef)	// Parent
							+ GetTableIndexSize(Table.Event);	// EventList
						break;
					case Table.EventPtr:
						size = GetTableIndexSize(Table.Event);	// Event
						break;
					case Table.Event:
						size = 2	// Flags
							+ stridx_size // Name
							+ GetCodedIndexSize(CodedIndex.TypeDefOrRef);	// EventType
						break;
					case Table.PropertyMap:
						size = GetTableIndexSize(Table.TypeDef)	// Parent
							+ GetTableIndexSize(Table.Property);	// PropertyList
						break;
					case Table.PropertyPtr:
						size = GetTableIndexSize(Table.Property);	// Property
						break;
					case Table.Property:
						size = 2	// Flags
							+ stridx_size	// Name
							+ blobidx_size;	// Type
						break;
					case Table.MethodSemantics:
						size = 2	// Semantics
							+ GetTableIndexSize(Table.Method)	// Method
							+ GetCodedIndexSize(CodedIndex.HasSemantics);	// Association
						break;
					case Table.MethodImpl:
						size = GetTableIndexSize(Table.TypeDef)	// Class
							+ GetCodedIndexSize(CodedIndex.MethodDefOrRef)	// MethodBody
							+ GetCodedIndexSize(CodedIndex.MethodDefOrRef);	// MethodDeclaration
						break;
					case Table.ModuleRef:
						size = stridx_size;	// Name
						break;
					case Table.TypeSpec:
						size = blobidx_size;	// Signature
						break;
					case Table.ImplMap:
						size = 2	// MappingFlags
							+ GetCodedIndexSize(CodedIndex.MemberForwarded)	// MemberForwarded
							+ stridx_size	// ImportName
							+ GetTableIndexSize(Table.ModuleRef);	// ImportScope
						break;
					case Table.FieldRVA:
						size = 4	// RVA
							+ GetTableIndexSize(Table.Field);	// Field
						break;
					case Table.EncLog:
					case Table.EncMap:
						size = 4;
						break;
					case Table.Assembly:
						size = 16 // HashAlgId 4, Version 4 * 2, Flags 4
							+ blobidx_size	// PublicKey
							+ (stridx_size * 2);	// Name, Culture
						break;
					case Table.AssemblyProcessor:
						size = 4;	// Processor
						break;
					case Table.AssemblyOS:
						size = 12;	// Platform 4, Version 2 * 4
						break;
					case Table.AssemblyRef:
						size = 12	// Version 2 * 4 + Flags 4
							+ (blobidx_size * 2)	// PublicKeyOrToken, HashValue
							+ (stridx_size * 2);	// Name, Culture
						break;
					case Table.AssemblyRefProcessor:
						size = 4	// Processor
							+ GetTableIndexSize(Table.AssemblyRef);	// AssemblyRef
						break;
					case Table.AssemblyRefOS:
						size = 12	// Platform 4, Version 2 * 4
							+ GetTableIndexSize(Table.AssemblyRef);	// AssemblyRef
						break;
					case Table.File:
						size = 4	// Flags
							+ stridx_size	// Name
							+ blobidx_size;	// HashValue
						break;
					case Table.ExportedType:
						size = 8	// Flags 4, TypeDefId 4
							+ (stridx_size * 2)	// Name, Namespace
							+ GetCodedIndexSize(CodedIndex.Implementation);	// Implementation
						break;
					case Table.ManifestResource:
						size = 8	// Offset, Flags
							+ stridx_size	// Name
							+ GetCodedIndexSize(CodedIndex.Implementation);	// Implementation
						break;
					case Table.NestedClass:
						size = GetTableIndexSize(Table.TypeDef)	// NestedClass
							+ GetTableIndexSize(Table.TypeDef);	// EnclosingClass
						break;
					case Table.GenericParam:
						size = 4	// Number, Flags
							+ GetCodedIndexSize(CodedIndex.TypeOrMethodDef)	// Owner
							+ stridx_size;	// Name
						break;
					case Table.MethodSpec:
						size = GetCodedIndexSize(CodedIndex.MethodDefOrRef)	// Method
							+ blobidx_size;	// Instantiation
						break;
					case Table.GenericParamConstraint:
						size = GetTableIndexSize(Table.GenericParam)	// Owner
							+ GetCodedIndexSize(CodedIndex.TypeDefOrRef);	// Constraint
						break;
					default:
						throw new NotSupportedException();
				}

				tables[i].RowSize = (uint)size;
				tables[i].Offset = offset;

				offset += (uint)size * tables[i].Length;
			}
		}

		public static Image ReadImageFrom(Stream stream)
		{
			try
			{
				var reader = new ImageReader(stream);
				reader.ReadImage();
				return reader.image;
			}
			catch(EndOfStreamException e)
			{
				throw new BadImageFormatException(stream.GetFullyQualifiedName(), e);
			}
		}
	}
}
